package com.example.graph;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class MST_Kruskal_Prim {
    public static void main(String[] args) {
        /**
         * {0, 6, 1, 5, 0, 0},
         * {6, 0, 5, 0, 3, 0},
         * {1, 5, 0, 5, 6, 4},
         * {5, 0, 5, 0, 0, 2},
         * {0, 3, 6, 0, 0, 6},
         * {0, 0, 4, 2, 6, 0}
         */
        int[][] edges;
        int n = 5;
        List<int[]> mstEdges;

        edges = new int[][] {{0, 1, 1}, {1, 2, 1}, {2, 3, 2}, {0, 3, 2}, {0, 4, 3}, {3, 4, 3}, {1, 4, 6}};
//        edges = new int[][]{
//                {0, 1, 6}, {0, 2, 1}, {0, 3, 5},
//                {1, 2, 5}, {1, 4, 3},
//                {2, 3, 5}, {2, 4, 6}, {2, 5, 4},
//                {3, 5, 2}, {4, 5, 6}};


        mstEdges = kruskal(n, edges);
        System.out.println("MST 权值:" + mstWeightSum);
        mstEdges.forEach(edge -> {
            System.out.println("[" + edge[0] + ", " + edge[1] + ", " + edge[2] + "]");
        });

        mstEdges = primByAdjMatrix(n, edges, 0);
//        mstEdges = primByAdjList(n, edges, 0);
        System.out.println("MST 权值:" + mstWeightSum);
        mstEdges.forEach(edge -> {
            System.out.println("[" + edge[0] + ", " + edge[1] + ", " + edge[2] + "]");
        });

    }

    /**
     * Kruskal算法查找MST
     * <p>
     * 算法思想:(贪心算法)
     * 以边为着眼点
     * 遍历所有边
     * 每次找最小的边
     * 判断当前最小边的两个顶点是否连通，(使用并查集完成判断)
     * 如果已经连通则跳过该边，否则将该边加入
     *
     * @return
     */

    private static int mstWeightSum = 0;

    public static List<int[]> kruskal(int n, int[][] edges) {
        List<int[]> res = new ArrayList<>();
        if (n == 0 || edges.length == 0 || edges[0].length == 0)
            return res;

        // 将边的数组按照权值从小到大排序
        Arrays.sort(edges, (e1, e2) -> e1[2] - e2[2]);
        // 并查集初始化
        init(n);

        for (int[] edge : edges) {
            if (union(edge[0], edge[1])) { // 如果需要合并说明两点不连通
                mstWeightSum += edge[2];
                res.add(edge);
            }
        }

        return res;
    }

    /**
     * Prim算法查找MST
     * <p>
     * 算法思想:(贪心算法)
     * 以顶点为着眼点
     * 首先初始化邻接表/矩阵,
     * 任选一个节点加入集合A，遍历A中元素找到最小权值边加入MST,并将该边的另一顶点加入A
     * 循环上述操作，累加最小权值
     *
     * @param n
     * @param edges
     * @param startNode
     * @return
     */
    public static List<int[]> primByAdjMatrix(int n, int[][] edges, int startNode) {
        List<int[]> res = new ArrayList<>();

        int[][] adjMatrix = initAdjMatrix(n, edges);
//        for (int i = 0; i < adjMatrix.length; i++) {
//            for (int j = 0; j < adjMatrix[0].length; j++) {
//                System.out.print(adjMatrix[i][j] + ", ");
//            }
//            System.out.println();
//        }

        // 初始化工作
        mstWeightSum = 0;
        int mstEdgeCount = 0;
        boolean[] visited = new boolean[n]; // 为true的节点表示已经合入集合A中,false表示该顶点还未加入A集合
        visited[startNode] = true; // 将指定开始节点加入A
        mstEdgeCount++;

        while (mstEdgeCount < n) {
            int x = -1, y = -1;
            int minWeight = Integer.MAX_VALUE;

            // 遍历A找到A中顶点里的最小边
            for (int i = 0; i < n; i++) {
                if (visited[i]) { // 只需要已被标记的元素(为true的元素表示在集合A中)
                    int[] neighbors = adjMatrix[i];
                    for (int j = 0; j < n; j++) { // 遍历邻居节点
                        if (neighbors[j] == 0) continue; // 两个顶点不连通
                        if (visited[j]) continue;
                        /* 记录较小权值 */
                        if (neighbors[j] < minWeight) {
                            minWeight = neighbors[j]; // 暂时记录当前最小权重和对应的边
                            x = i;y = j; // 暂时记录边的两个顶点
                        }
                    }
                }
            }

            visited[y] = true; // 将最小边的另一个顶点加入到集合A中
            mstEdgeCount++;
            mstWeightSum += minWeight; // 累计权重
            res.add(new int[]{x, y, adjMatrix[x][y]}); // 将边加入MST中
        }
        return res;
    }

    /**
     * MST 权值:15
     * [0, 2, 1]
     * [2, 5, 4]
     * [5, 3, 2]
     * [2, 1, 5]
     * [1, 4, 3]
     *
     * @param n
     * @param edges
     * @param startNode
     * @return
     */

    public static List<int[]> primByAdjList(int n, int[][] edges, int startNode) {
        List<int[]> res = new ArrayList<>();
        if (n == 0 || edges.length == 0 || edges[0].length == 0)
            return res;

        List<List<int[]>> adj = initAdjList(n, edges);
//        for (int i = 0; i < adj.size(); i++) {
//            System.out.print(i + "--->");
//            adj.get(i).forEach(e -> System.out.print("[" + e[0] + "," + e[1] + "]"));
//            System.out.println();
//        }

        // 初始化工作
        mstWeightSum = 0;
        int mstEdgeCount = 0;
        boolean[] visited = new boolean[n]; // 为true的节点表示已经合入集合A中,false表示该顶点还未加入A集合
        visited[startNode] = true; // 将指定开始节点加入A
        mstEdgeCount++;

        while (mstEdgeCount < n) {
            int x = -1, y = -1;
            int minWeight = Integer.MAX_VALUE;

            for (int i = 0; i < n; i++) {
                if (visited[i]) {
                    List<int[]> neighbors = adj.get(i);
                    for (int j = 0; j < neighbors.size(); j++) {
                        int[] neighbor = neighbors.get(j);
                        if (!visited[neighbor[0]] && neighbor[1] < minWeight) { // 邻居节点未加入集合A并且边的权值比当前小
                            minWeight = neighbor[1];
                            x = i; y = neighbor[0];
                        }
                    }
                }
            }

            if (x != -1 || y != -1) {
                visited[y] = true; // 将最小边的另一个顶点加入到集合A中
            }
            mstEdgeCount++;
            mstWeightSum += minWeight; // 累计权重
            res.add(new int[]{x, y, minWeight}); // 将边加入MST中

        }

        return res;
    }


    private static List<List<int[]>> initAdjList(int n, int[][] edges) {
        List<List<int[]>> adj = new ArrayList<>();
        for (int i = 0; i < n; i++) {
            adj.add(new ArrayList<>());
        }

        for (int[] edge : edges) {
            adj.get(edge[0]).add(new int[]{edge[1], edge[2]});
            adj.get(edge[1]).add(new int[]{edge[0], edge[2]});
        }

        return adj;
    }

    private static int[][] initAdjMatrix(int n, int[][] edges) {
        int[][] matrix = new int[n][n];

        for (int[] edge : edges) {
            matrix[edge[0]][edge[1]] = edge[2];
            matrix[edge[1]][edge[0]] = edge[2];
        }
        return matrix;
    }

    /*******************并查集操作******************/
    private static int[] parent;
    private static int[] rank;

    // 初始化并查集
    private static void init(int n) {
        parent = new int[n]; // 存每个节点的父节点
        rank = new int[n];// 存每个节点的秩

        for (int i = 0; i < n; i++) {
            parent[i] = i;
            rank[i] = 1;
        }

    }

    // 查找指定节点的父节点,同时完成路径压缩
    private static int find(int node) {
        while (parent[node] != node) {
            node = parent[node];
            parent[node] = parent[parent[node]];
        }

        return parent[node];
    }

    // 合并两个集合,不需要合并返回false,需要合并的话,按秩合并后返回true
    private static boolean union(int set1, int set2) {
        int root1 = find(set1);
        int root2 = find(set2);
        if (root1 != root2) {
            if (rank[root1] > rank[root2]) {
                parent[root2] = root1;

            } else if (rank[root1] < rank[root2]) {
                parent[root1] = root2;
            } else {
                parent[root2] = root1;
                rank[root1] += 1;
            }
            return true;
        }
        return false;
    }
    /*********************并查集操作end**************************/

}
